import torch
import torch.nn as nn

# 准备数据     binary_cross_entropy函数要求输入和目标都是Float类型
batch_size = 5
torch.manual_seed(42)
x = torch.arange(5).unsqueeze(1).float()
label_name = {
    0: "猫",
    1: "狗",
    2: "狸花猫",
    3: "简州猫",
    4: "金毛犬",
    5: "哈士奇"
}
# 单标签二分类：只有一个标签，是或不是（乳腺癌的预测）
# 多标签二分类：多个标签，每个标签下面依旧是或不是，也就是0或1
y = torch.tensor([
    [1, 0, 0, 1, 0, 0],
    [0, 1, 0, 0, 1, 0],
    [1, 0, 1, 0, 0, 0],
    [0, 1, 0, 0, 0, 1],
    [1, 0, 0, 0, 0, 0]
]).float()

# 1. 超参数
lr = 0.01
Epochs = 10000

# 2. 定义模型
model = nn.Sequential(
    nn.Linear(1, 64),
    nn.ReLU(),
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 6),
    nn.Sigmoid()
)

# 3. 损失函数
# nn.BCEWithLogitsLoss()      自带Sigmoid激活
loss_fn = nn.BCELoss()

# 4. 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 5. 循环训练
model.train()
# model.train() 是PyTorch中将模型设置为训练模式的方法
# 启用训练模式后，模型会启用如Dropout、BatchNorm等训练时特有的行为
# 与model.eval()相对，后者用于验证/测试阶段
for epoch in range(Epochs):
    # 前向传播
    y_pre = model(x)
    # 计算损失
    loss = loss_fn(y_pre, y)
    # 清空梯度
    optimizer.zero_grad()
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()

    # 打印训练信息
    if epoch == 0 or (epoch + 1) % 100 == 0:
        print(f"[{epoch + 1}/{Epochs}] Loss:{loss.item():.4f}")


# ------------------------------------------------------------
def pre_label_name(x):
    model.eval()
    # 进入推理模式的时候，我们只要输出，是不需要梯度信息的，为了节省内存，提升计算效率
    # 可以关闭梯度计算
    with torch.no_grad():
        y_pre = model(x)
        output = torch.round(y_pre)
        ret = []
        for out in output:
            label = [label_name[i] for i, val in enumerate(out) if val == 1]
            ret.append(", ".join(label))
        return ret


pres = pre_label_name(x)
for i, pre in enumerate(pres):
    print(f"第{i + 1}张图片的分类是：{pre}")
